// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <ktest/ktest.h>

#include <nebula/core/array.h>
#include <nebula/array/builder_decimal.h>
#include <nebula/core/datum.h>
#include <nebula/core/record_batch.h>
#include <nebula/types/tensor.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/testing/random.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>

namespace nebula {

// Test basic cases for contains NaN.
class TestAssertContainsNaN : public ::testing::Test {};

TEST_F(TestAssertContainsNaN, BatchesEqual) {
  auto schema = ::nebula::schema({
      {field("a", float32())},
      {field("b", float64())},
  });

  auto expected = assert_record_batch_from_json(schema,
                                      R"([{"a": 3,    "b": 5},
                                       {"a": 1,    "b": 3},
                                       {"a": 3,    "b": 4},
                                       {"a": NaN,  "b": 6},
                                       {"a": 2,    "b": 5},
                                       {"a": 1,    "b": NaN},
                                       {"a": 1,    "b": 3}
                                       ])");
  auto actual = assert_record_batch_from_json(schema,
                                    R"([{"a": 3,    "b": 5},
                                       {"a": 1,    "b": 3},
                                       {"a": 3,    "b": 4},
                                       {"a": NaN,  "b": 6},
                                       {"a": 2,    "b": 5},
                                       {"a": 1,    "b": NaN},
                                       {"a": 1,    "b": 3}
                                       ])");
  ASSERT_BATCHES_EQUAL(*expected, *actual);
  AssertBatchesApproxEqual(*expected, *actual);
}

TEST_F(TestAssertContainsNaN, TableEqual) {
  auto schema = ::nebula::schema({
      {field("a", float32())},
      {field("b", float64())},
  });

  auto expected = assert_table_from_json(schema, {R"([{"a": null, "b": 5},
                                     {"a": NaN,    "b": 3},
                                     {"a": 3,    "b": null}
                                    ])",
                                         R"([{"a": null, "b": null},
                                     {"a": 2,    "b": NaN},
                                     {"a": 1,    "b": 5},
                                     {"a": 3,    "b": 5}
                                    ])"});
  auto actual = assert_table_from_json(schema, {R"([{"a": null, "b": 5},
                                     {"a": NaN,    "b": 3},
                                     {"a": 3,    "b": null}
                                    ])",
                                       R"([{"a": null, "b": null},
                                     {"a": 2,    "b": NaN},
                                     {"a": 1,    "b": 5},
                                     {"a": 3,    "b": 5}
                                    ])"});
  ASSERT_TABLES_EQUAL(*expected, *actual);
}

TEST_F(TestAssertContainsNaN, ArrayEqual) {
  auto expected = assert_array_from_json(float64(), "[0, 1, 2, NaN]");
  auto actual = assert_array_from_json(float64(), "[0, 1, 2, NaN]");
  AssertArraysEqual(*expected, *actual);
}

TEST_F(TestAssertContainsNaN, ChunkedEqual) {
  auto expected = assert_chunked_array_from_json(float64(), {
                                                      "[null, 1]",
                                                      "[3, NaN, 2]",
                                                      "[NaN]",
                                                  });

  auto actual = assert_chunked_array_from_json(float64(), {
                                                    "[null, 1]",
                                                    "[3, NaN, 2]",
                                                    "[NaN]",
                                                });
  AssertChunkedEqual(*expected, *actual);
}

TEST_F(TestAssertContainsNaN, DatumEqual) {
  // scalar
  auto expected_scalar = assert_scalar_from_json(float64(), "NaN");
  auto actual_scalar = assert_scalar_from_json(float64(), "NaN");
  AssertDatumsEqual(expected_scalar, actual_scalar);

  // array
  auto expected_array = assert_array_from_json(float64(), "[3, NaN, 2, 1, 5]");
  auto actual_array = assert_array_from_json(float64(), "[3, NaN, 2, 1, 5]");
  AssertDatumsEqual(expected_array, actual_array);

  // chunked array
  auto expected_chunked = assert_chunked_array_from_json(float64(), {
                                                              "[null, 1]",
                                                              "[3, NaN, 2]",
                                                              "[NaN]",
                                                          });

  auto actual_chunked = assert_chunked_array_from_json(float64(), {
                                                            "[null, 1]",
                                                            "[3, NaN, 2]",
                                                            "[NaN]",
                                                        });
  AssertDatumsEqual(expected_chunked, actual_chunked);
}

class TestTensorFromJSON : public ::testing::Test {};

TEST_F(TestTensorFromJSON, FromJSONAndArray) {
  std::vector<int64_t> shape = {9, 2};
  const int64_t i64_size = sizeof(int64_t);
  std::vector<int64_t> f_strides = {i64_size, i64_size * shape[0]};
  std::vector<int64_t> f_values = {1,  2,  3,  4,  5,  6,  7,  8,  9,
                                   10, 20, 30, 40, 50, 60, 70, 80, 90};
  auto data = Buffer::Wrap(f_values);

  std::shared_ptr<Tensor> tensor_expected;
  ASSERT_OK_AND_ASSIGN(tensor_expected, Tensor::create(int64(), data, shape, f_strides));

  std::shared_ptr<Tensor> result = assert_tensor_from_json(
      int64(), "[1, 2,  3,  4,  5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90]",
      shape, f_strides);

  EXPECT_TRUE(tensor_expected->equals(*result));
}

TEST_F(TestTensorFromJSON, FromJSON) {
  std::vector<int64_t> shape = {9, 2};
  std::vector<int64_t> values = {1,  2,  3,  4,  5,  6,  7,  8,  9,
                                 10, 20, 30, 40, 50, 60, 70, 80, 90};
  auto data = Buffer::Wrap(values);

  std::shared_ptr<Tensor> tensor_expected;
  ASSERT_OK_AND_ASSIGN(tensor_expected, Tensor::create(int64(), data, shape));

  std::shared_ptr<Tensor> result = assert_tensor_from_json(
      int64(), "[1, 2,  3,  4,  5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60, 70, 80, 90]",
      "[9, 2]");

  EXPECT_TRUE(tensor_expected->equals(*result));
}

}  // namespace nebula
